# blive_danmu.py
# -------------
# Author: Zako888y
#
# This file is dedicated to the public domain under the Creative Commons CC0 1.0 Universal (CC0 1.0) Public Domain Dedication.
# You may use, copy, modify, distribute, and/or sell this code for any purpose, without restriction. No warranty is provided. The author waives all rights and claims to this work worldwide.
# To view a copy of this dedication, visit https://creativecommons.org/publicdomain/zero/1.0/
#
# 此文件参考自 原作者：鱼肉烧 https://www.bilibili.com/opus/595528208079862415 的方案。
# 原代码的主逻辑由 JavaScript 编写，参见 reference/live_bili_danmu.html。
# 原代码被完全仿照并改写为 Python 版，两个文件间不存在运行上的依赖关系。
#

import asyncio
import json
from re import S
import struct
from abc import ABC, abstractmethod

import brotli
import websockets


class OnDanmuListener(ABC):
    """弹幕监听回调接口"""

    __slots__ = ()

    @abstractmethod
    def onDanmu(self, uname: str, text: str) -> None:
        """弹幕回调"""
        pass

    def onPopularity(self, popularity: int) -> None:
        """人气值回调，默认不处理"""
        pass


class DanmuClient:

    _running: bool
    _ws: websockets.ClientConnection | None
    _heartbeat_task: asyncio.Task | None
    _main_task: asyncio.Task | None

    def __init__(
        self, uid: int, roomid: int, token: str, callback: OnDanmuListener
    ) -> None:
        self._uid = uid
        self._roomid = roomid
        self._token = token
        self._callback = callback

        self._running = False
        self._ws = None
        self._heartbeat_task = None
        self._main_task = None

    @property
    def uid(self) -> int:
        """访客 UID"""
        return self._uid

    @property
    def roomid(self) -> int:
        """主播直播间 ID"""
        return self._roomid

    @property
    def token(self) -> str:
        """访客 Token"""
        return self._token

    @property
    def callback(self) -> OnDanmuListener:
        """弹幕事件"""
        return self._callback

    async def start(self) -> None:
        """启动弹幕监听"""
        self._running = True
        certification = {
            "uid": self.uid,
            "roomid": self.roomid,
            "protover": 3,
            "platform": "web",
            "type": 2,
            "key": self.token,
        }
        try:
            async with websockets.connect(self.BLIVE_CHAT_URI) as ws:
                self._ws = ws
                await ws.send(self._get_certification(certification))
                print("认证包已发送")

                async def heartbeat():
                    while self._running and self._ws:
                        try:
                            await self._ws.send(self._get_heartbeat())
                        except Exception:
                            pass
                        await asyncio.sleep(30)

                self._heartbeat_task = asyncio.create_task(heartbeat())
                while self._running:
                    try:
                        msg = await ws.recv()
                        if isinstance(msg, bytes):
                            self._handle_message(msg)
                        else:
                            print("服务器消息：", msg)
                    except websockets.exceptions.ConnectionClosedError as e:
                        print(f"验证失败：{e}")
                        break
                    except Exception as e:
                        print(f"未知错误：{e}")
                        break
        except Exception as e:
            print(f"连接异常：{e}")
        finally:
            self._running = False
            if self._heartbeat_task:
                self._heartbeat_task.cancel()
            self._ws = None

    def stop(self) -> None:
        """停止弹幕监听"""
        self._running = False
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
        if self._main_task:
            self._main_task.cancel()

    def _handle_message(self, data: bytes) -> None:
        """处理服务器数据"""
        offset = 0
        while offset < len(data):
            packet_len = struct.unpack_from(">I", data, offset)[0]
            head_len = struct.unpack_from(">H", data, offset + 4)[0]
            ver = struct.unpack_from(">H", data, offset + 6)[0]
            op = struct.unpack_from(">I", data, offset + 8)[0]
            body = data[offset + head_len : offset + packet_len]
            if ver == 3:
                body = brotli.decompress(body)
                self._handle_message(body)  # 递归处理解压后的数据
            else:
                if op != 3:
                    try:
                        msg = json.loads(body.decode("utf-8"))
                        # 只处理弹幕事件 DANMU_MSG
                        if isinstance(msg, dict) and msg.get("cmd") == "DANMU_MSG":
                            info = msg.get("info", [])
                            # info[2][1] 是用户名，info[1] 是内容
                            uname = (
                                info[2][1] if len(info) > 2 and len(info[2]) > 1 else ""
                            )
                            text = info[1] if len(info) > 1 else ""
                            self.callback.onDanmu(uname, text)
                    except Exception:
                        pass

                else:
                    # 人气值；虽然是没用的功能但还是保留吧
                    popularity = struct.unpack_from(">I", body)[0]
                    self.callback.onPopularity(popularity)

            offset += packet_len

    BLIVE_CHAT_URI: str = "wss://broadcastlv.chat.bilibili.com:443/sub"

    @staticmethod
    def _get_certification(data: dict) -> bytes:
        """生成认证包"""
        json_bytes = json.dumps(data).encode("utf-8")
        packet_len = len(json_bytes) + 16
        header = struct.pack(">IHHII", packet_len, 16, 1, 7, 1)
        return header + json_bytes

    @staticmethod
    def _get_heartbeat() -> bytes:
        """生成心跳包"""
        return struct.pack(">IHHII", 16, 16, 1, 2, 1)


if __name__ == "__main__":
    import sys

    if len(sys.argv) == 4:

        class DanmuPrinter(OnDanmuListener):
            def __init__(self, show_popularity: bool = False):
                self.show_popularity = show_popularity

            def onDanmu(self, uname: str, text: str) -> None:
                print(f"{uname}: {text}")

            def onPopularity(self, popularity: int) -> None:
                if self.show_popularity:
                    print(f"人气值: {popularity}")

        uid = int(sys.argv[1])
        roomid = int(sys.argv[2])
        token = sys.argv[3]
        callback = DanmuPrinter()
        client = DanmuClient(uid, roomid, token, callback)

        async def run_client():
            try:
                await client.start()
            except KeyboardInterrupt:
                print("已停止")
                client.stop()
            finally:
                client.stop()

        asyncio.run(run_client())
    else:
        print(
            """用法: python blive_danmu.py <uid> <roomid> <token>
            uid: 你的哔哩哔哩 UID
            roomid: 直播间号
            token: 通过登录获取的 token

            示例: python blive_danmu.py 2093276158 22849290 b-nxcCUG3M_rCnPgczgM_p9tfiKK4WvzdFvxrZj2NY36wMeFS9-bMvf-ZQ5WCAhUWGKGoaqXWw-OWTiGgSZm_BSDKFJ1oJEeLY_Tes_IuJVYwvPmE-Aox6-qcVzSJOOMwfMTUIhCg2C5yf9-kaaIvKGfEpmDFnmXly-lWBK_0dETDu_UGo5Usl2L-nOo7YovhBb-5i-yZIwDAhVM_2TnosGk6cipaLlm5FRItZBjVHUDw5iG3DpZ5Qg0m6o0KMYUscp946Nwh0zCw8ZwQEAqObTs
            """
        )
